#pragma once
#include <string_view>
namespace mms {
    class BitStream {
    public:
        BitStream(const std::string_view& stream) : stream_(stream) {

        }

        bool read_one_bit(uint8_t& b) {
            if (bit_offset_ >= 8) {
                byte_index_++;
                bit_offset_ = 0;
            }
            uint8_t v;
            try {
                v = (uint8_t)stream_.at(byte_index_);
            }
            catch (std::exception& e) {
                return false;
            }

            b = (v >> (7 - bit_offset_)) & 0x01;
            bit_offset_++;
            return true;
        }

        bool write_one_bit(uint8_t v) {
            if (bit_offset_ >= 8) {
                byte_index_++;
                bit_offset_ = 0;
            }

            char* byte = (char*)stream_.data() + byte_index_;
            if (v == 1) {
                *byte = *byte | (v << (7 - bit_offset_));
            }
            else {
                *byte = *byte & ~(v << (7 - bit_offset_));
            }
            bit_offset_++;
            return true;
        }

        bool read_bits(size_t count, uint64_t& v) {
            v = 0;
            for (size_t i = 0; i < count; i++) {
                v = v << 1;
                uint8_t b;
                if (!read_one_bit(b)) {
                    return false;
                }
                v = v | b;
            }
            return true;
        }

        bool write_bits(size_t count, uint64_t v) {
            for (size_t i = 0; i < count; i++) {
                uint8_t b = (v >> (count - i - 1)) & 0x01;
                if (!write_one_bit(b)) {
                    return false;
                }
                // v = v >> 1;
            }
            return true;
        }

        bool read_bits(size_t count, uint8_t& v) {
            v = 0;
            for (size_t i = 0; i < count; i++) {
                v = v << 1;
                uint8_t b;
                if (!read_one_bit(b)) {
                    return false;
                }
                v = v | b;
            }
            return true;
        }

        bool read_bits(size_t count, uint16_t& v) {
            v = 0;
            for (size_t i = 0; i < count; i++) {
                v = v << 1;
                uint8_t b;
                if (!read_one_bit(b)) {
                    return false;
                }
                v = v | b;
            }
            return true;
        }

        /*
         * Parsing process for Exp-Golomb codes
            This process is invoked when the descriptor of a syntax element in the syntax tables in clause 7.3 is equal to ue(v), me(v),
            se(v), or te(v). For syntax elements in clauses 7.3.4 and 7.3.5, this process is invoked only when
            entropy_coding_mode_flag is equal to 0.
            Inputs to this process are bits from the RBSP.
            Outputs of this process are syntax element values.
            Syntax elements coded as ue(v), me(v), or se(v) are Exp-Golomb-coded. Syntax elements coded as te(v) are truncated ExpGolomb-coded. The parsing process for these syntax elements begins with reading the bits starting at the current location
            in the bitstream up to and including the first non-zero bit, and counting the number of leading bits that are equal to 0. This
            process is specified as follows:
            leadingZeroBits = −1
            for( b = 0; !b; leadingZeroBits++ ) (9-1)
            b = read_bits( 1 )
            The variable codeNum is then assigned as follows:
            codeNum = 2^leadingZeroBits − 1 + read_bits( leadingZeroBits ) (9-2)
        */
        bool read_exp_golomb_ue(uint64_t& v) {
            int leadingZeroBits = -1;
            uint8_t b = 0;
            for (b = 0; !b && leadingZeroBits < 64; leadingZeroBits++) {
                if (!read_one_bit(b)) {
                    return false;
                }
            }
            uint64_t l;
            if (!read_bits(leadingZeroBits, l)) {
                return false;
            }
            v = (1 << leadingZeroBits) - 1 + l;
            return true;
        }

        bool read_exp_golomb_se(int64_t& v) {
            uint64_t ue;
            if (!read_exp_golomb_ue(ue)) {
                return false;
            }
            bool positive = ue & 0x01;
            v = (ue + 1) >> 1;
            if (!positive) {
                v = -v;
            }
            return true;
        }
    public:
        std::string_view stream_;
        size_t byte_index_ = 0;
        size_t bit_offset_ = 0;
    };
};